#!/usr/bin/env python
import sys
import signal
from killerbee import *

def usage():
    print >>sys.stderr, """
zbdump - a tcpdump-like tool for ZigBee/IEEE 802.15.4 networks
Compatible with Wireshark 1.1.2 and later - jwright@willhackforsushi.com

Usage: zbdump [-fiwDch] [-f channel] [-w pcapfile] [-W daintreefile] 
         [-i devnumstring]
    """

def show_dev():
    kb = KillerBee()
    print "Dev\tProduct String\tSerial Number"
    for dev in kb.dev_list():
        print "%s\t%s\t%s" % (dev[0], dev[1], dev[2])

def interrupt(signum, frame):
    global packetcount
    global kb
    global pd
    global dt
    kb.sniffer_off()
    kb.close()
    if pd:
        pd.close()
    if dt:
        dt.close()
    print "%d packets captured" % packetcount
    sys.exit(0)

# PcapDumper, only used if -w is specified
pd = None
# DainTreeDumper, only used if -W is specified
dt = None

# Command-line arguments
arg_channel = None
arg_pcapsavefile = None
arg_dsnasavefile = None
arg_devstring = None
arg_count = -1

# Global
packetcount = 0

while len(sys.argv) > 1:
    op = sys.argv.pop(1)
    if op == '-f':
        arg_channel = int(sys.argv.pop(1))
    if op == '-w':
        arg_pcapsavefile = sys.argv.pop(1)
    if op == '-W':
        arg_dsnasavefile = sys.argv.pop(1)
    if op == '-i':
        arg_devstring = sys.argv.pop(1)
    if op == '-h':
        usage()
        sys.exit(0)
    if op == '-D':
        show_dev()
        sys.exit(0)
    if op == '-c':
        arg_count = int(sys.argv.pop(1))
 
if arg_channel == None:
    print >>sys.stderr, "ERROR: Must specify a channel with -f"
    usage()
    sys.exit(1)

if arg_pcapsavefile == None and arg_dsnasavefile == None:
    print >>sys.stderr, "ERROR: Must specify a savefile with -w (libpcap) or -W (Daintree SNA)"
    usage()
    sys.exit(1)

if (arg_pcapsavefile != None):
    pd = PcapDumper(DLT_IEEE802_15_4, arg_pcapsavefile)
if (arg_dsnasavefile != None):
    dt = DainTreeDumper(arg_dsnasavefile)

kb = KillerBee(device=arg_devstring)
signal.signal(signal.SIGINT, interrupt)

kb.set_channel(arg_channel)
kb.sniffer_on()

print "zbdump: listening on \'%s\', link-type DLT_IEEE802_15_4, capture size 127 bytes" % kb.get_dev_info()[0]

while arg_count != packetcount:
    packet = kb.pnext()
    # packet[1] is True if CRC is correct
    if packet != None and packet[1]:
        packetcount+=1
        if pd:
            # Adding a "\x00" causes FCS's to be bad, thanks Tyler!
            #packet[0] = packet[0][:-1] + "\x00"
            pd.pcap_dump(packet[0][:-1])
        if dt:
            dt.pwrite(packet[0])

kb.sniffer_off()
kb.close()
if pd:
    pd.close()
if dt:
    dt.close()

print "%d packets captured" % packetcount
